'''
Created on Apr 4, 2012

@author: pedro
'''
from Utils import *
from random import Random,randint
from Neurons import spikes_to_epsp, alpha_synapse, rc_synapse, stdp
from functools import reduce
from matplotlib import pyplot as plt
from math import sin, pi

import Pulses
from CustomLoggers import MyLogger

class OutputLayer(object):
    
    def __init__(self, pyramidal_cells, output_neurons, stdp_decay_factor, max_weight, connectivity):
        self.pyramidal_cells = pyramidal_cells
        self.output_neurons = output_neurons
        self.stdp_decay_factor = stdp_decay_factor
        self.max_weight = max_weight
        
        '''
        pyr_to_out_weights is in the form of {(Px,Oy):weight}, where weight comes from
        the connectivityMatrix calculated in the appropriate method.
        '''
        self.pyr_to_out_weights = {}
        
        connectivityMatrix = self.initialize_connectivity_matrix(connectivity)
        for x in range(len(connectivityMatrix)):
            for y in range(len(connectivityMatrix[x])):
                self.pyr_to_out_weights[(self.pyramidal_cells[y],self.output_neurons[x])] = connectivityMatrix[x][y]
        return
    
    def initialize_connectivity_matrix(self, percent_subset):
        '''
        Each PyramidalNeuron connects to a %-subset of all OutputNeurons
        '''
        MyLogger.my_logger.info("Initializing connectivity matrix...")

        connectivity_matrix = [[0.00 for val2 in range(len(self.output_neurons))] for val in range(len(self.pyramidal_cells))]
        
        for x,row in enumerate(connectivity_matrix):
            #Pick random subset to connect to
            indices = []
            while len(indices) < len(row):
                next_index = randint(0,len(row)-1)
                if next_index in indices:
                    continue
                indices.append(next_index)  
            indices = indices[:int(round(percent_subset*len(indices)))]
            for y in range(len(row)):
                row[y] = 1.00 if y in indices else 0.00
            
        #DEBUG
        MyLogger.my_logger.debug("Connectivity Matrix for pyr_to_out_weights:")
        for row in connectivity_matrix:
            MyLogger.my_logger.debug(row)
                
        MyLogger.my_logger.info("... Connectivity matrix initialized")

        return connectivity_matrix
    
    
    def reset_neurons(self):
        for pyr in self.pyramidal_cells:
            pyr.reset_vm()
        for out in self.output_neurons:
            out.reset_vm()
        
    
    def fire(self, pyramidal_outputs, dt, time, learning, Aplus, Aminus, stdp_tauplus, stdp_tauminus):
        
        MyLogger.my_logger.info("Started network firing seqeunce...")

        pyramidal_outputs_dict = {}
        for pyramidal_output in pyramidal_outputs:
            pyramidal_outputs_dict[pyramidal_output[0]] = pyramidal_output[1]
        
        '''Step 0: Generate one combined EPSP per pyramidal neuron and map to a dictionary, i.e. {pyr : epsp}'''
        
        MyLogger.my_logger.info("Generating and combining Pyramidal-to-Output EPSPs")
        input_pulses = {}
        
        for pyramidal_output in pyramidal_outputs:
            epsp = spikes_to_epsp(pyramidal_output[1],alpha_synapse,time,dt) if pyramidal_output[1][0][0] != 0.00 else [0.00 for val in range(time)]
            input_pulses[pyramidal_output[0]]= epsp
        
        MyLogger.my_logger.info("Pyramidal-to-Output EPSPs Generated")

        
        '''Step 1: Fire the Output Neurons and mark their firing times'''
        
        MyLogger.my_logger.info("Firing Output Neurons")

        output_firing_times = {}
        output_timecourses = []
        
        for output_neuron in self.output_neurons:
            input_pyrs = []
            for pyr in self.pyramidal_cells:
                weight = self.pyr_to_out_weights[(pyr,output_neuron)]
                input_pyrs.append(map(lambda x:x*weight,input_pulses[pyr]))
            pulse = reduce(list_accumulate,input_pyrs)
            output_firing_times[output_neuron] = output_neuron.time_of_fires(pulse,dt)
            output_neuron.reset_vm()
            output_timecourses.append((output_neuron,output_neuron.fire(pulse,dt)))
        
        #DEBUG
        MyLogger.my_logger.debug("Output Firing Times: ")
        for o in output_firing_times.keys():
            MyLogger.my_logger.debug("({0},{1})".format(o,output_firing_times[o]))
        
        MyLogger.my_logger.info("Output neurons fired")
        


        '''Step 2: Apply STDP Learning Rule to Alter Synaptic Weights'''

        MyLogger.my_logger.info("Applying STDP Learning Rule")

        for pyr in self.pyramidal_cells:
            
            #If learning is off, then skip this step:
            if (learning is False):
                #DEBUG:
                MyLogger.my_logger.info("Learning is off; synaptic weights unchanged")
                break;
            
            for output_neuron in self.output_neurons:
                
                weight = self.pyr_to_out_weights[(pyr,output_neuron)]
                out_fires = output_firing_times[output_neuron]
                
                #If OutputNeuron didn't fire, then don't do STDP
                if weight == 0.00 or out_fires[0][0] == -1.0:
                    continue
                
                for item in out_fires:
                    time_fired = item[0]
                    
                    #Check if Pyramidal neuron fired: if not, decrease weight by self.stdp_decay_factor
                    if pyramidal_outputs_dict[pyr][0][0] == -1.0:
                        weight = weight * 1/(1+self.stdp_decay_factor)
                        MyLogger.my_logger.debug("{0} never fired, so weight is now {1}".format(pyr,weight))
                        continue
                    
                    time_deltas = [val[0] - time_fired for val in pyramidal_outputs_dict[pyr]]
                    positive_weight_change = []
                    negative_weight_change = []
                    
                    for delta_t in time_deltas:
                        if delta_t < 0:
                            positive_weight_change.append(stdp(delta_t,stdp_tauplus,Aplus,Aminus)+1.0)
                        else:
                            negative_weight_change.append(stdp(delta_t,stdp_tauminus,Aplus,Aminus)+1.0)
                            
                    positive_weight_change = [1.00] if len(positive_weight_change) == 0 or pyramidal_outputs_dict[pyr][0] == -1.0 else positive_weight_change
                    negative_weight_change = [1.00 + self.stdp_decay_factor] if len(negative_weight_change) == 0 else negative_weight_change
                    
                    total_weight_change = sum(positive_weight_change)/(self.stdp_decay_factor + sum(negative_weight_change))
                    
                    weight = weight * total_weight_change
                    
                    #If weight is outside acceptable bounds, sigmoid it
                    if weight > self.max_weight:
                        weight = sigmoidal_weight_filter(weight, self.max_weight)
                    
                        
                    MyLogger.my_logger.debug("Time Deltas: {0}".format(time_deltas))
                    MyLogger.my_logger.debug("Positive Weight Change for Pyr {0} and Out {1} is {2}".format(pyr,output_neuron,sum(positive_weight_change)))
                    MyLogger.my_logger.debug("Negative Weight Change for Pyr {0} and Out {1} is {2}".format(pyr,output_neuron,sum(negative_weight_change)))
                    MyLogger.my_logger.debug("TOTAL Weight Change for Pyr {0} and Out {1} is {2}".format(pyr,output_neuron,total_weight_change))
                    MyLogger.my_logger.info("NEW Weight for Pyr {0} and Out {1} is {2}".format(pyr,output_neuron,weight))
                    
                    
                self.pyr_to_out_weights[(pyr,output_neuron)] = weight
        
        log_list_as_matrix(self.pyr_to_out_weights.values(),len(self.pyramidal_cells),len(self.output_neurons))
        
        MyLogger.my_logger.info("... Network Firing seqeunce completed")

        return output_timecourses
 

'''-------------------------------------------------------------------------------------------'''


'''Uses 40 Hz Oscillator'''
class OscillatorPyramidalLayer(object):
    
    def __init__(self, pyramidal_cells, oscillator_freq):
        self.pyramidal_cells = pyramidal_cells
        self.oscillator_freq = oscillator_freq #In Hertz (Hz)
        return

    def reset_neurons(self):
        for pyr in self.pyramidal_cells:
            pyr.reset_vm()
        return
    
    def oscillator(self, t):
        return 1000 * sin((2.0*pi*1.0/self.oscillator_freq)*(t+(self.oscillator_freq/2.0)))
    
    ''''Firing Method'''
    def fire(self, hdr, dt, time):

        MyLogger.my_logger.info("Firing Oscillator Pyramidal Layer...")

        '''
        Step #0: Simulate all RC pulses
        '''
        pulses = []
        for x in range(len(hdr)):
            rc_pulse = Pulses.RCFilteredPulse(hdr[x]*100.0,time,10.0,7.0)
            pulses.append([next(rc_pulse) for t in range(time)])
        
        '''
        Step #1: Add oscillator timecourse to all rc_pulses
        '''
        oscillator_timecourse = []
        
        for x in range(time):
            oscillator_timecourse.append(self.oscillator(x))
            
        for pulse_index in range(len(pulses)):
            pulses[pulse_index] = [val + osc for val,osc in zip(pulses[pulse_index],oscillator_timecourse)]
        
        '''
        Step #2: Put into pyramidal_inputs in the form of (Py, timecourse)
        '''
        pyramidal_inputs = []
        for y in range(len(self.pyramidal_cells)):
            pyramidal_inputs.append((self.pyramidal_cells[y], pulses[y]))
        
        '''
        Step #3: Prepare for pyramidal firing
        '''        
        pyramidal_times_of_firing = []
        pyramidal_timecourses = []
        list_of_pyrs = self.pyramidal_cells
        
        '''
        Step #4: Simulate all the Pyramidal Cells
        '''
        for item in pyramidal_inputs:
            index = list_of_pyrs.index(item[0])
            modulated_input = pulses[index]
            pyramidal_times_of_firing.append( (item[0], item[0].time_of_fires( modulated_input, dt)) )
            item[0].reset_vm()
            pyramidal_timecourses.append( (item[0],item[0].fire(modulated_input, dt)))
        
        MyLogger.my_logger.info("... OscillatorPyramidalLayer fired")

        '''
        Step #4b: (Optional) Plot Firing Patterns of Pyramidal Cells
        '''
        '''
        for pyr, timecourse in pyramidal_timecourses:
            plt.plot(timecourse, label=str(pyr))
        plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=6, ncol=2, mode="expand", borderaxespad=0.)
        plt.title("Pyramidal Cell Firing Timecourse")
        plt.show()
        '''
        
        return pyramidal_times_of_firing


'''-------------------------------------------------------------------------------------------'''

class PILayer(object):
    '''
    classdocs
    '''

    def __init__(self, pyramidal_cells, interneurons):
        self.pyramidal_cells = pyramidal_cells
        self.interneurons = interneurons
        
        numInterneuronsPerPyr = len(interneurons)/len(pyramidal_cells)
        chnks = chunks(interneurons, numInterneuronsPerPyr)
        self.columns = [ (pyramidal_cells[x], chnks[x]) for x in range(len(pyramidal_cells)) ]
        
        '''
        i_to_pyr_weights is in the form of {(Ix,Py):weight}, where weight comes from
        the connectivityMatrix calculated in the appropriate method.
        '''
        self.i_to_pyr_weights = {}
        
        connectivityMatrix = self.initialize_connectivity_matrix(0.20,0.20,len(pyramidal_cells),numInterneuronsPerPyr)
        for x in range(len(connectivityMatrix)):
            for y in range(len(connectivityMatrix[x])):
                self.i_to_pyr_weights[(self.interneurons[y],self.pyramidal_cells[x])] = connectivityMatrix[x][y]
        return
    
    
    def initialize_connectivity_matrix(self, minConnectivity, maxConnectivity, numPyrs, numInterneuronsPerPyr):
        r = Random(1231231)
        connectivityMatrix = [ range(numPyrs*numInterneuronsPerPyr) for x in range(numPyrs)]
        
        for x in range(len(connectivityMatrix)):
            currPyrConnectivity = minConnectivity + (r.random() * (maxConnectivity - minConnectivity))
            numAssociatedInterneurons = int(round( currPyrConnectivity* (len(connectivityMatrix[x]) - numInterneuronsPerPyr) ))
            associatedInterneurons = []
            
            while (numAssociatedInterneurons > 0):
                indexAssociated = r.randint(0, len(connectivityMatrix[x]))
                
                if (indexAssociated >= x*numInterneuronsPerPyr) and (indexAssociated < (x*numInterneuronsPerPyr)+numInterneuronsPerPyr):
                    pass
                elif ((indexAssociated in associatedInterneurons) == False):
                    associatedInterneurons.append(indexAssociated)
                    numAssociatedInterneurons = numAssociatedInterneurons - 1
            
            for y in range(len(connectivityMatrix[x])):
                if (y in associatedInterneurons):
                    connectivityMatrix[x][y] = 1.0
                else:
                    connectivityMatrix[x][y] = 0.0
                    
        return connectivityMatrix
    
    
    def reset_neurons(self):
        for pyr in self.pyramidal_cells:
            pyr.reset_vm(-0.07)
        for inter in self.interneurons:
            inter.reset_vm(-0.07)
        return
   
   
    ''''Firing Method'''
    def fire(self, hdr, dt, time):
        '''
        Step #0: Simulate all RC pulses
        '''
        pulses = []
        for x in range(len(hdr)):
            rc_pulse = Pulses.RCFilteredPulse(hdr[x]*100.0,500.0,10.0,7.0)
            pulses.append([next(rc_pulse) for t in range(time)])
        
        '''
        Step #1: Fire all the interneurons with appropriate columnar input
        '''    
        interneuron_ouputs = []
        for x in range(len(self.columns)):
            for i in self.columns[x][1]:
                interneuron_ouputs.append(i.fire(pulses[x],dt))
        
        pyramidal_interneuron_inputs = []
        
        '''
        Step #2: Multiply each interneuron ouput Vm timecourse by its synaptic weight
        to each pyramidal cell, put into pyramidal_interneuron_inputs in the form of (Py, weight*[timecourse])
        '''
        for x in range(len(self.interneurons)):
            for y in range(len(self.pyramidal_cells)):
                synaptic_weight = self.i_to_pyr_weights[(self.interneurons[x],self.pyramidal_cells[y])]
                pyramidal_interneuron_inputs.append((self.pyramidal_cells[y],[val*synaptic_weight for val in interneuron_ouputs[x]]))
        
        '''
        Step #3: Accumulate duplicate time courses in preparation for mitral firing
        '''
        pyramidal_interneuron_inputs = accumulate_duplicates(pyramidal_interneuron_inputs)
        
        pyramidal_times_of_firing = []
        pyramidal_timecourses = []
        
        list_of_pyrs = [val[0] for val in self.columns]
        
        '''
        Step #4: Simulate all the Pyramidal Cells
        '''
        for item in pyramidal_interneuron_inputs:
            index = list_of_pyrs.index(item[0])
            #print "UnModulated Input = {0}".format(pulses[index])
            modulated_input = list_subtract(pulses[index],item[1])
            #print "Modulated Input = {0}".format(modulated_input)
            pyramidal_times_of_firing.append( (item[0], item[0].time_of_fires( modulated_input, dt)) )
            item[0].reset_vm(-0.070)
            pyramidal_timecourses.append( (item[0],item[0].fire(modulated_input, dt)))
        
        '''
        Step #4b: (Optional) Plot Firing Patterns of Pyramidal Cells
        
        for pyr, timecourse in pyramidal_timecourses:
            plt.plot(timecourse, label=str(pyr))
        plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=6, ncol=2, mode="expand", borderaxespad=0.)
        plt.title("Pyramidal Cell Firing Timecourse")
        plt.show()
        '''
        return pyramidal_times_of_firing
        